package com.pig4cloud.pigx.knowledge.support.rule;

import cn.hutool.core.util.StrUtil;
import com.pig4cloud.pigx.common.core.util.R;
import com.pig4cloud.pigx.knowledge.dto.AiMessageResultDTO;
import com.pig4cloud.pigx.knowledge.dto.ChatMessageDTO;
import com.pig4cloud.pigx.knowledge.service.AiFunctionAssistantService;
import com.pig4cloud.pigx.knowledge.support.feign.RemoteTableInfoService;
import com.pig4cloud.pigx.knowledge.support.function.Chat2SqlFunctionCalling;
import com.pig4cloud.pigx.knowledge.support.provider.ModelProvider;
import com.pig4cloud.pigx.knowledge.support.util.ChatMessageContextHolder;
import com.pig4cloud.pigx.knowledge.support.util.PromptBuilder;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.Map;

/**
 * 基于数据库的聊天
 *
 * @author lengleng
 * @date 2024/3/26
 */
@Slf4j
@Component("databaseChat")
@RequiredArgsConstructor
public class DatabaseChatRule implements ChatRule {

	private final Chat2SqlFunctionCalling chat2SqlFunctionCalling;

	private final RemoteTableInfoService tableInfoService;

	private final ModelProvider modelProvider;

	/**
	 * 处理聊天信息
	 * @param chatMessageDTO 聊天上文
	 * @return flux stream
	 */
	@Override
	public Flux<AiMessageResultDTO> process(ChatMessageDTO chatMessageDTO) {
		ChatMessageDTO.Table tableInfo = chatMessageDTO.getExtDetails().getTable();
		StringBuilder tableSchemas = new StringBuilder();

		for (String tableName : Arrays.stream(tableInfo.getTableNames()).distinct().toList()) {
			R<String> columnR = tableInfoService.listColumn(tableInfo.getDsName(), tableName);
			tableSchemas.append(columnR.getData()).append(StrUtil.CRLF);
		}

		ChatMessageDTO.ExtDetails extDetails = new ChatMessageDTO.ExtDetails();
		extDetails.setFuncName(chat2SqlFunctionCalling.functionName());
		chatMessageDTO.setExtDetails(extDetails);
		ChatMessageContextHolder.set(chatMessageDTO);

		AiFunctionAssistantService functionAssistantService = modelProvider
			.getAiFunctionAssistant(chatMessageDTO.getModelName())
			.getValue();
		String result = functionAssistantService.chat(chatMessageDTO.getConversationId(),
				PromptBuilder.render("chat2db.st",
						Map.of("tableSchema", tableSchemas.toString(), "userInput", chatMessageDTO.getContent())));
		return Flux.just(new AiMessageResultDTO(result));
	}

}
